#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
化学反应计算引擎
支持化学反应平衡、热力学计算、物料平衡等
"""

import re
import json
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass, asdict
from pathlib import Path
import logging

logger = logging.getLogger(__name__)


@dataclass
class Reactant:
    """反应物/产物数据结构"""
    formula: str
    coefficient: float
    molecular_weight: float
    name: str = ""
    phase: str = ""  # s, l, g, aq
    
    def to_dict(self) -> Dict:
        return asdict(self)


@dataclass
class ReactionThermodynamics:
    """反应热力学数据"""
    delta_H: Optional[float] = None  # 焓变 kJ/mol
    delta_G: Optional[float] = None  # 吉布斯自由能变 kJ/mol
    delta_S: Optional[float] = None  # 熵变 J/(mol·K)
    equilibrium_constant: Optional[float] = None  # 平衡常数
    temperature: float = 298.15  # K
    
    def to_dict(self) -> Dict:
        return asdict(self)


@dataclass
class ChemicalReaction:
    """化学反应数据结构"""
    reaction_id: str
    reaction_name: str
    reaction_equation: str
    reactants: List[Reactant]
    products: List[Reactant]
    thermodynamics: Optional[ReactionThermodynamics] = None
    reaction_type: str = ""
    catalyst: Optional[str] = None
    conditions: Dict = None
    
    def to_dict(self) -> Dict:
        data = {
            'reaction_id': self.reaction_id,
            'reaction_name': self.reaction_name,
            'reaction_equation': self.reaction_equation,
            'reactants': [r.to_dict() for r in self.reactants],
            'products': [p.to_dict() for p in self.products],
            'reaction_type': self.reaction_type,
            'catalyst': self.catalyst,
            'conditions': self.conditions or {}
        }
        if self.thermodynamics:
            data['thermodynamics'] = self.thermodynamics.to_dict()
        return data


class ReactionParser:
    """化学反应方程式解析器"""
    
    @staticmethod
    def parse_equation(equation: str) -> Tuple[List[Dict], List[Dict]]:
        """
        解析化学方程式
        
        Args:
            equation: 化学方程式，如 "2H2 + O2 → 2H2O"
            
        Returns:
            (反应物列表, 产物列表)
        """
        # 分离反应物和产物
        if '→' in equation:
            left, right = equation.split('→')
        elif '->' in equation:
            left, right = equation.split('->')
        elif '=' in equation:
            left, right = equation.split('=')
        else:
            raise ValueError(f"无法解析反应方程式: {equation}")
        
        reactants = ReactionParser._parse_side(left.strip())
        products = ReactionParser._parse_side(right.strip())
        
        return reactants, products
    
    @staticmethod
    def _parse_side(side: str) -> List[Dict]:
        """解析方程式一侧"""
        compounds = []
        
        # 分离各个化合物
        parts = side.split('+')
        
        for part in parts:
            part = part.strip()
            
            # 提取系数
            match = re.match(r'^(\d*\.?\d*)\s*(.+)$', part)
            if match:
                coef_str, formula = match.groups()
                coefficient = float(coef_str) if coef_str else 1.0
            else:
                coefficient = 1.0
                formula = part
            
            # 提取相态
            phase = ""
            if '(' in formula:
                formula_match = re.match(r'(.+)\(([slgaq]+)\)', formula)
                if formula_match:
                    formula, phase = formula_match.groups()
            
            compounds.append({
                'formula': formula.strip(),
                'coefficient': coefficient,
                'phase': phase
            })
        
        return compounds


class ReactionCalculator:
    """化学反应计算器"""
    
    # 常见元素原子量
    ATOMIC_WEIGHTS = {
        'H': 1.008, 'C': 12.011, 'N': 14.007, 'O': 15.999,
        'S': 32.065, 'P': 30.974, 'Cl': 35.453, 'Br': 79.904,
        'F': 18.998, 'I': 126.904, 'Na': 22.990, 'K': 39.098,
        'Ca': 40.078, 'Mg': 24.305, 'Fe': 55.845, 'Cu': 63.546,
        'Zn': 65.38, 'Al': 26.982, 'Si': 28.085, 'Ag': 107.868,
        'Au': 196.967, 'Pt': 195.084, 'Ni': 58.693, 'Mn': 54.938,
    }
    
    # 气体常数
    R = 8.314  # J/(mol·K)
    
    def __init__(self):
        """初始化计算器"""
        pass
    
    @staticmethod
    def calculate_molecular_weight(formula: str) -> float:
        """
        计算分子量
        
        Args:
            formula: 分子式，如 "H2O", "C6H12O6"
            
        Returns:
            分子量 (g/mol)
        """
        # 解析分子式
        pattern = r'([A-Z][a-z]?)(\d*)'
        matches = re.findall(pattern, formula)
        
        total_weight = 0.0
        
        for element, count in matches:
            count = int(count) if count else 1
            
            if element in ReactionCalculator.ATOMIC_WEIGHTS:
                total_weight += ReactionCalculator.ATOMIC_WEIGHTS[element] * count
            else:
                logger.warning(f"未知元素: {element}")
        
        return round(total_weight, 3)
    
    @staticmethod
    def balance_mass(reaction: ChemicalReaction, 
                     reactant_amounts: Dict[str, float]) -> Dict[str, float]:
        """
        物料平衡计算
        
        Args:
            reaction: 化学反应对象
            reactant_amounts: 反应物质量 {formula: mass_in_kg}
            
        Returns:
            产物质量 {formula: mass_in_kg}
        """
        # 找到限制性反应物
        limiting_reactant = None
        min_moles = float('inf')
        
        for reactant in reaction.reactants:
            if reactant.formula in reactant_amounts:
                mass = reactant_amounts[reactant.formula]
                moles = mass * 1000 / reactant.molecular_weight  # kg转g
                normalized_moles = moles / reactant.coefficient
                
                if normalized_moles < min_moles:
                    min_moles = normalized_moles
                    limiting_reactant = reactant
        
        if not limiting_reactant:
            return {}
        
        # 计算产物质量
        product_masses = {}
        
        for product in reaction.products:
            product_moles = min_moles * product.coefficient
            product_mass = (product_moles * product.molecular_weight) / 1000  # g转kg
            product_masses[product.formula] = round(product_mass, 3)
        
        return product_masses
    
    @staticmethod
    def calculate_yield(theoretical_mass: float, 
                       actual_mass: float) -> float:
        """
        计算收率
        
        Args:
            theoretical_mass: 理论产量 (kg)
            actual_mass: 实际产量 (kg)
            
        Returns:
            收率 (0-1)
        """
        if theoretical_mass == 0:
            return 0.0
        return round(actual_mass / theoretical_mass, 4)
    
    @staticmethod
    def calculate_equilibrium_constant(delta_G: float, 
                                      temperature: float = 298.15) -> float:
        """
        根据吉布斯自由能计算平衡常数
        
        Args:
            delta_G: 吉布斯自由能变 (kJ/mol)
            temperature: 温度 (K)
            
        Returns:
            平衡常数 K
        """
        # ΔG = -RT ln(K)
        # K = exp(-ΔG / RT)
        delta_G_J = delta_G * 1000  # kJ转J
        K = pow(2.71828, -delta_G_J / (ReactionCalculator.R * temperature))
        return K
    
    @staticmethod
    def calculate_gibbs_energy(delta_H: float, 
                              delta_S: float, 
                              temperature: float = 298.15) -> float:
        """
        计算吉布斯自由能
        
        Args:
            delta_H: 焓变 (kJ/mol)
            delta_S: 熵变 (J/(mol·K))
            temperature: 温度 (K)
            
        Returns:
            吉布斯自由能 (kJ/mol)
        """
        # ΔG = ΔH - TΔS
        delta_G = delta_H - (temperature * delta_S / 1000)
        return round(delta_G, 2)
    
    @staticmethod
    def is_spontaneous(delta_G: float) -> bool:
        """
        判断反应是否自发进行
        
        Args:
            delta_G: 吉布斯自由能变 (kJ/mol)
            
        Returns:
            True表示自发，False表示非自发
        """
        return delta_G < 0
    
    @staticmethod
    def calculate_heat_released(delta_H: float, 
                               product_moles: float) -> float:
        """
        计算反应放热量
        
        Args:
            delta_H: 焓变 (kJ/mol)
            product_moles: 产物摩尔数
            
        Returns:
            放热量 (kJ)，负值表示放热
        """
        return delta_H * product_moles


class ReactionDesigner:
    """化学反应设计器"""
    
    def __init__(self, compound_db=None):
        """
        初始化反应设计器
        
        Args:
            compound_db: 化合物数据库实例
        """
        self.compound_db = compound_db
        self.parser = ReactionParser()
        self.calculator = ReactionCalculator()
    
    def create_reaction(self, 
                       reaction_id: str,
                       reaction_name: str,
                       equation: str,
                       compound_data: Dict[str, Dict] = None,
                       thermodynamics: Dict = None,
                       **kwargs) -> ChemicalReaction:
        """
        创建化学反应对象
        
        Args:
            reaction_id: 反应ID
            reaction_name: 反应名称
            equation: 化学方程式
            compound_data: 化合物数据 {formula: {name, mw, ...}}
            thermodynamics: 热力学数据
            **kwargs: 其他参数
            
        Returns:
            化学反应对象
        """
        # 解析方程式
        reactant_data, product_data = self.parser.parse_equation(equation)
        
        # 创建反应物对象
        reactants = []
        for r in reactant_data:
            formula = r['formula']
            mw = self.calculator.calculate_molecular_weight(formula)
            
            name = formula
            if compound_data and formula in compound_data:
                name = compound_data[formula].get('name', formula)
                mw = compound_data[formula].get('molecular_weight', mw)
            
            reactants.append(Reactant(
                formula=formula,
                coefficient=r['coefficient'],
                molecular_weight=mw,
                name=name,
                phase=r.get('phase', '')
            ))
        
        # 创建产物对象
        products = []
        for p in product_data:
            formula = p['formula']
            mw = self.calculator.calculate_molecular_weight(formula)
            
            name = formula
            if compound_data and formula in compound_data:
                name = compound_data[formula].get('name', formula)
                mw = compound_data[formula].get('molecular_weight', mw)
            
            products.append(Reactant(
                formula=formula,
                coefficient=p['coefficient'],
                molecular_weight=mw,
                name=name,
                phase=p.get('phase', '')
            ))
        
        # 创建热力学对象
        thermo = None
        if thermodynamics:
            thermo = ReactionThermodynamics(**thermodynamics)
        
        # 创建反应对象
        reaction = ChemicalReaction(
            reaction_id=reaction_id,
            reaction_name=reaction_name,
            reaction_equation=equation,
            reactants=reactants,
            products=products,
            thermodynamics=thermo,
            **kwargs
        )
        
        return reaction
    
    def export_reaction(self, reaction: ChemicalReaction, output_file: str):
        """导出反应到JSON文件"""
        output_path = Path(output_file)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(reaction.to_dict(), f, ensure_ascii=False, indent=2)
        
        logger.info(f"反应已导出到: {output_file}")

